//
// The MIT License (MIT)
//
// Copyright (c) 2022 Livox. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
//

#define _USE_MATH_DEFINES
#include <cmath>

#include "livox_lidar_def.h"
#include "livox_lidar_api.h"

#ifdef WIN32
#include <windows.h>
#else
#include <unistd.h>
#endif

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <thread>
#include <chrono>
#include <iostream>
#include <string>
#include <vector>
#include <mutex>
#include <condition_variable>
#include <map>
#include <fstream>
#include <sstream>
#include <iomanip>

#include <csignal>

// 如果M_PI仍未定义，手动定义
#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

#ifdef _WIN32
#else
#include <sys/stat.h>
#include <sys/types.h>
#include <errno.h>
#endif

std::condition_variable quit_condition;
std::mutex mtx;

// 点云数据结构
struct PointCloudData {
  float x, y, z;
  uint8_t intensity;
  uint64_t timestamp;
};

// 全局变量
std::vector<PointCloudData> point_cloud_buffer;
std::mutex point_cloud_mutex;
uint32_t total_points_saved = 0;
uint32_t file_counter = 0;
const uint32_t POINTS_PER_FILE = 100000;  // 每个文件保存的点数（不再触发中间保存，仅统计用）

// 保存目录
const char* kDefaultRootSaveDir = "/home/unitree/softwares/Livox-SDK2-master/pd_files";
std::string g_session_time_str;
std::string g_save_dir;

// 函数声明
bool SavePointsToPLY(const std::vector<PointCloudData>& points, const std::string& filename);
bool SavePointsToPCD(const std::vector<PointCloudData>& points, const std::string& filename);
std::string GetCurrentTimeString();
void ProcessPointCloudData(uint32_t handle, const uint8_t dev_type, LivoxLidarEthernetPacket* data);
uint64_t ConvertTimestamp(const uint8_t timestamp[8]);
bool MakeDirsRecursively(const std::string& path);

// 递归创建目录（仅UNIX路径样式）
bool MakeDirsRecursively(const std::string& path) {
#ifdef _WIN32
  // Windows 简化处理：一次性创建（可根据需要扩展）
  return _mkdir(path.c_str()) == 0 || errno == EEXIST;
#else
  if (path.empty()) return false;
  if (path[0] != '/') return false; // 只处理绝对路径

  std::string partial;
  partial.reserve(path.size());
  for (size_t i = 0; i < path.size(); ++i) {
    partial.push_back(path[i]);
    if (path[i] == '/' && !partial.empty()) {
      if (partial.size() == 1) continue; // root '/'
      if (mkdir(partial.c_str(), 0755) != 0 && errno != EEXIST) {
        // 继续尝试创建后续目录可能也会失败，这里直接返回
        // printf("mkdir failed: %s (%d)\n", partial.c_str(), errno);
        return false;
      }
    }
  }
  // 最后一段如果不是'/'结尾，仍需创建
  if (partial.back() != '/') {
    if (mkdir(partial.c_str(), 0755) != 0 && errno != EEXIST) {
      return false;
    }
  }
  return true;
#endif
}

// 将timestamp数组转换为uint64_t
uint64_t ConvertTimestamp(const uint8_t timestamp[8]) {
  uint64_t result = 0;
  for (int i = 0; i < 8; i++) {
    result |= ((uint64_t)timestamp[i]) << (i * 8);
  }
  return result;
}

// 获取当前时间字符串
std::string GetCurrentTimeString() {
  auto now = std::chrono::system_clock::now();
  auto time_t = std::chrono::system_clock::to_time_t(now);
  std::stringstream ss;
  ss << std::put_time(std::localtime(&time_t), "%Y%m%d_%H%M%S");
  return ss.str();
}

// 保存点云数据到PLY文件
bool SavePointsToPLY(const std::vector<PointCloudData>& points, const std::string& filename) {
  std::ofstream ply_file(filename);
  if (!ply_file.is_open()) {
    printf("Error: Cannot open PLY file: %s\n", filename.c_str());
    return false;
  }

  // 写入PLY文件头
  ply_file << "ply\n";
  ply_file << "format ascii 1.0\n";
  ply_file << "comment Created by Livox SDK2\n";
  ply_file << "element vertex " << points.size() << "\n";
  ply_file << "property float x\n";
  ply_file << "property float y\n";
  ply_file << "property float z\n";
  ply_file << "property uchar intensity\n";
  ply_file << "property double timestamp\n";
  ply_file << "end_header\n";

  // 写入点云数据
  for (const auto& point : points) {
    ply_file << std::fixed << std::setprecision(6)
             << point.x << " " << point.y << " " << point.z << " "
             << (int)point.intensity << " " << point.timestamp << "\n";
  }

  ply_file.close();
  return true;
}

// 保存点云数据到PCD文件
bool SavePointsToPCD(const std::vector<PointCloudData>& points, const std::string& filename) {
  std::ofstream pcd_file(filename);
  if (!pcd_file.is_open()) {
    printf("Error: Cannot open PCD file: %s\n", filename.c_str());
    return false;
  }

  // 写入PCD文件头
  pcd_file << "# .PCD v0.7 - Point Cloud Data file format\n";
  pcd_file << "VERSION 0.7\n";
  pcd_file << "FIELDS x y z intensity timestamp\n";
  pcd_file << "SIZE 4 4 4 1 8\n";
  pcd_file << "TYPE F F F U U\n";
  pcd_file << "COUNT 1 1 1 1 1\n";
  pcd_file << "WIDTH " << points.size() << "\n";
  pcd_file << "HEIGHT 1\n";
  pcd_file << "VIEWPOINT 0 0 0 1 0 0 0\n";
  pcd_file << "POINTS " << points.size() << "\n";
  pcd_file << "DATA ascii\n";

  // 写入点云数据
  for (const auto& point : points) {
    pcd_file << std::fixed << std::setprecision(6)
             << point.x << " " << point.y << " " << point.z << " "
             << (int)point.intensity << " " << point.timestamp << "\n";
  }

  pcd_file.close();
  return true;
}

// 处理点云数据
void ProcessPointCloudData(uint32_t handle, const uint8_t dev_type, LivoxLidarEthernetPacket* data) {
  if (data == nullptr) {
    return;
  }

  std::vector<PointCloudData> current_points;
  uint64_t timestamp = ConvertTimestamp(data->timestamp);
  
  // 解析不同格式的点云数据
  if (data->data_type == kLivoxLidarCartesianCoordinateHighData) {
    LivoxLidarCartesianHighRawPoint *points = (LivoxLidarCartesianHighRawPoint *)data->data;
    for (uint32_t i = 0; i < data->dot_num; i++) {
      PointCloudData point;
      point.x = points[i].x * 0.001f;  // 转换为米
      point.y = points[i].y * 0.001f;
      point.z = points[i].z * 0.001f;
      point.intensity = points[i].reflectivity;
      point.timestamp = timestamp;
      current_points.push_back(point);
    }
  } else if (data->data_type == kLivoxLidarCartesianCoordinateLowData) {
    LivoxLidarCartesianLowRawPoint *points = (LivoxLidarCartesianLowRawPoint *)data->data;
    for (uint32_t i = 0; i < data->dot_num; i++) {
      PointCloudData point;
      point.x = points[i].x * 0.01f;   // 转换为米 (cm to m)
      point.y = points[i].y * 0.01f;
      point.z = points[i].z * 0.01f;
      point.intensity = points[i].reflectivity;
      point.timestamp = timestamp;
      current_points.push_back(point);
    }
  } else if (data->data_type == kLivoxLidarSphericalCoordinateData) {
    LivoxLidarSpherPoint *points = (LivoxLidarSpherPoint *)data->data;
    for (uint32_t i = 0; i < data->dot_num; i++) {
      PointCloudData point;
      // 球坐标转直角坐标
      float r = points[i].depth * 0.001f;  // 转换为米
      float theta = points[i].theta * 0.01f * M_PI / 180.0f;  // 转换为弧度
      float phi = points[i].phi * 0.01f * M_PI / 180.0f;
      
      point.x = r * std::sin(theta) * std::cos(phi);
      point.y = r * std::sin(theta) * std::sin(phi);
      point.z = r * std::cos(theta);
      point.intensity = points[i].reflectivity;
      point.timestamp = timestamp;
      current_points.push_back(point);
    }
  }

  // 添加到缓冲区（不再中途保存，统一在退出时保存为单个文件）
  {
    std::lock_guard<std::mutex> lock(point_cloud_mutex);
    point_cloud_buffer.insert(point_cloud_buffer.end(), current_points.begin(), current_points.end());
    total_points_saved = static_cast<uint32_t>(point_cloud_buffer.size());
  }
}

void PointCloudCallback(uint32_t handle, const uint8_t dev_type, LivoxLidarEthernetPacket* data, void* client_data) {
  if (data == nullptr) {
    return;
  }
  
  static uint32_t frame_count = 0;
  frame_count++;
  
  // 每100帧打印一次状态
  if (frame_count % 100 == 0) {
    printf("Received frame %u: handle=%u, data_num=%d, data_type=%d, length=%d\n",
           frame_count, handle, data->dot_num, data->data_type, data->length);
  }
  
  // 处理点云数据
  ProcessPointCloudData(handle, dev_type, data);
}

void ImuDataCallback(uint32_t handle, const uint8_t dev_type, LivoxLidarEthernetPacket* data, void* client_data) {
  if (data == nullptr) {
    return;
  }
  static uint32_t imu_count = 0;
  imu_count++;
  
  // 每50帧打印一次IMU状态
  if (imu_count % 50 == 0) {
    printf("IMU data: handle=%u, frame=%u, data_num=%u, data_type=%u\n",
           handle, imu_count, data->dot_num, data->data_type);
  }
}

void WorkModeCallback(livox_status status, uint32_t handle, LivoxLidarAsyncControlResponse *response, void *client_data) {
  if (response == nullptr) {
    return;
  }
  printf("WorkMode Callback: status=%u, handle=%u, ret_code=%u, error_key=%u\n",
         status, handle, response->ret_code, response->error_key);
}

void LoggerStartCallback(livox_status status, uint32_t handle, LivoxLidarLoggerResponse* response, void* client_data) {
  if (status != kLivoxLidarStatusSuccess) {
    printf("Start logger failed, status: %d\n", status);
    return;
  }

  if (response == nullptr) {
    printf("Start logger failed, response is nullptr\n");
    return;
  }

  if (response->ret_code != 0) {
    printf("Start logger failed, ret_code: %d\n", response->ret_code);
    return;
  }

  printf("Lidar[%u] start logger successfully\n", handle);
}

void DebugPointCloudCallback(livox_status status, uint32_t handle, LivoxLidarLoggerResponse* response, void* client_data) {
  printf("Debug PointCloud Callback: status=%d, handle=%u", status, handle);
  if (response != nullptr) {
    printf(", ret_code=%d", response->ret_code);
  }
  printf("\n");
}

void LidarInfoChangeCallback(const uint32_t handle, const LivoxLidarInfo* info, void* client_data) {
  if (info == nullptr) {
    printf("Lidar info change callback failed, info is nullptr\n");
    return;
  }
  
  printf("Lidar Info Change: handle=%u, SN=%s\n", handle, info->sn);
  
  // 设置雷达工作模式为正常模式
  SetLivoxLidarWorkMode(handle, kLivoxLidarNormal, WorkModeCallback, nullptr);
  
  // 启动日志记录
  LivoxLidarStartLogger(handle, kLivoxLidarRealTimeLog, LoggerStartCallback, nullptr);
  
  // 启用调试点云
  SetLivoxLidarDebugPointCloud(handle, true, DebugPointCloudCallback, nullptr);
  
  printf("Lidar[%u] configuration completed, starting point cloud collection...\n", handle);
}

void Stop(int signal) {
  printf("\nReceived stop signal, saving merged PLY/PCD to session folder...\n");

  // 构建保存目录路径
  if (g_session_time_str.empty()) {
    g_session_time_str = GetCurrentTimeString();
  }
  g_save_dir = std::string(kDefaultRootSaveDir) + "/" + g_session_time_str;

  // 创建目录
  if (!MakeDirsRecursively(g_save_dir)) {
    printf("Error: Cannot create save directory: %s\n", g_save_dir.c_str());
  }

  // 构建文件名
  std::string ply_filename = g_save_dir + "/pointcloud_merged.ply";
  std::string pcd_filename = g_save_dir + "/pointcloud_merged.pcd";

  // 保存剩余的点云数据
  std::lock_guard<std::mutex> lock(point_cloud_mutex);
  if (!point_cloud_buffer.empty()) {
    if (SavePointsToPLY(point_cloud_buffer, ply_filename)) {
      printf("Saved merged PLY file: %s (%zu points)\n", ply_filename.c_str(), point_cloud_buffer.size());
    }
    if (SavePointsToPCD(point_cloud_buffer, pcd_filename)) {
      printf("Saved merged PCD file: %s (%zu points)\n", pcd_filename.c_str(), point_cloud_buffer.size());
    }
    total_points_saved = static_cast<uint32_t>(point_cloud_buffer.size());
    file_counter = 1;
  }
  
  printf("Total points saved: %u, Total files created: %u\n", total_points_saved, file_counter);
  quit_condition.notify_all();
}

int main(int argc, const char *argv[]) {
  if (argc != 2) {
    printf("Usage: %s <config_path>\n", argv[0]);
    printf("Example: %s mid360_config.json\n", argv[0]);
    return -1;
  }
  
  const std::string path = argv[1];
  g_session_time_str = GetCurrentTimeString();
  g_save_dir = std::string(kDefaultRootSaveDir) + "/" + g_session_time_str;

  // 预先创建保存目录
  if (MakeDirsRecursively(g_save_dir)) {
    printf("Save directory: %s\n", g_save_dir.c_str());
  } else {
    printf("Warning: Failed to create save directory: %s\n", g_save_dir.c_str());
  }

  printf("Starting Livox Point Cloud Collector...\n");
  printf("Config file: %s\n", path.c_str());
  printf("Points per file (stat only): %u\n", POINTS_PER_FILE);

  // 初始化SDK
  if (!LivoxLidarSdkInit(path.c_str())) {
    printf("Livox SDK initialization failed\n");
    LivoxLidarSdkUninit();
    return -1;
  }

  // 设置回调函数
  SetLivoxLidarPointCloudCallBack(PointCloudCallback, nullptr);
  SetLivoxLidarImuDataCallback(ImuDataCallback, nullptr);
  SetLivoxLidarInfoChangeCallback(LidarInfoChangeCallback, nullptr);

  // 注册信号处理
  std::signal(SIGINT, Stop);

  printf("Waiting for lidar connection...\n");
  printf("Press Ctrl+C to stop and save data\n");

  // 等待退出信号
  std::unique_lock<std::mutex> lock(mtx);
  quit_condition.wait(lock);

  printf("Shutting down Livox SDK...\n");
  LivoxLidarSdkUninit();
  printf("Point Cloud Collector finished!\n");
  
  return 0;
} 